import numpy as np
import h5py

# 定义文件路径
h5_file_path = '/home/crxc/disk/emReg/data/Mec_label0.2.h5'
new_h5_file_path = '/home/crxc/disk/emReg/data/Mec_label_cor0.2.h5'

# 打开原始 HDF5 文件和创建新的 HDF5 文件
with h5py.File(h5_file_path, 'r') as original_f, h5py.File(new_h5_file_path, 'w') as corrected_f:
    # 复制原始文件内容到新文件
    original_f.copy('image', corrected_f)

    # 获取原始 'image' 组中 'flow' 数据集的长度
    num_images = original_f['image']['flow'].shape[0]
    # 如果数据集已经存在，首先删除它
    if 'image/flow' in corrected_f:
        del corrected_f['image/flow']
    # 创建修正后的 'flow' 数据集
    corrected_flow_dataset = corrected_f.create_dataset('image/flow',
                                                        shape=(num_images, 512, 512, 2),
                                                        maxshape=(None, 512, 512, 2),
                                                        dtype=np.float32)

    # 遍历每个光流数据集，并进行修正
    for i in range(num_images):
        # 读取带有坐标偏移的光流数据
        flow_with_offset = original_f['image']['flow'][i, ...]

        # 减去坐标偏移
        corrected_flow = flow_with_offset.astype(np.float32)
        corrected_flow[:, :, 0] -= np.arange(512).astype(np.float32)
        corrected_flow[:, :, 1] -= np.arange(512).reshape(-1, 1).astype(np.float32)

        # 将修正后的光流数据保存到新的数据集中
        corrected_flow_dataset[i, ...] = corrected_flow
        print(i)
    # 如果还有其他数据需要修正（例如：flow_img, wrap_img），则重复上述步骤
    # ...

    # 确保新的 HDF5 文件数据被写入并且文件被正确关闭
    corrected_f.flush()
